import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import os
import matplotlib.pyplot as plt

# 设备选择
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
])
dataset = torchvision.datasets.MNIST(root="./data", train=True, transform=transform, download=True)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

# 生成器
class Generator(nn.Module):
    def __init__(self, z_dim, y_dim, img_dim):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(z_dim + y_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, img_dim),
            nn.Tanh()
        )

    def forward(self, z, y):
        inputs = torch.cat((z, y), dim=1)
        return self.model(inputs).view(-1, 1, 28, 28)

# 判别器
class Discriminator(nn.Module):
    def __init__(self, img_dim, y_dim):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(img_dim + y_dim, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 128),
            nn.LeakyReLU(0.2),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, x, y):
        inputs = torch.cat((x.view(x.size(0), -1), y), dim=1)
        return self.model(inputs)

# 初始化网络
z_dim = 100
img_dim = 784
y_dim = 10  # MNIST类别
generator = Generator(z_dim, y_dim, img_dim).to(device)
discriminator = Discriminator(img_dim, y_dim).to(device)

# 损失函数与优化器
criterion = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=0.0002)
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002)

# 训练CGAN
num_epochs = 150
fixed_noise = torch.randn(16, z_dim).to(device)
fixed_labels = torch.zeros(16, y_dim).to(device)
fixed_labels[:, 5] = 1  # 仅生成数字 5

for epoch in range(num_epochs):
    for images, labels in dataloader:
        labels = labels.to(device)  # 迁移标签到 GPU
        labels_one_hot = torch.zeros(labels.size(0), y_dim).to(device)
        labels_one_hot.scatter_(1, labels.view(-1, 1), 1)

        images = images.to(device)

        # 训练判别器
        real_labels = torch.ones(images.size(0), 1).to(device)
        fake_labels = torch.zeros(images.size(0), 1).to(device)

        optimizer_D.zero_grad()
        outputs = discriminator(images, labels_one_hot)
        loss_real = criterion(outputs, real_labels)

        noise = torch.randn(images.size(0), z_dim).to(device)
        fake_images = generator(noise, labels_one_hot)
        outputs = discriminator(fake_images.detach(), labels_one_hot)
        loss_fake = criterion(outputs, fake_labels)

        loss_D = loss_real + loss_fake
        loss_D.backward()
        optimizer_D.step()

        # 训练生成器
        optimizer_G.zero_grad()
        outputs = discriminator(fake_images, labels_one_hot)
        loss_G = criterion(outputs, real_labels)
        loss_G.backward()
        optimizer_G.step()

    # 记录损失值
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss D: {loss_D.item()}, Loss G: {loss_G.item()}")

    # 每隔固定代数生成图片
    if epoch % 10 == 0:
        generated_images = generator(fixed_noise, fixed_labels)
        torchvision.utils.save_image(generated_images, f"generated_epoch_{epoch}.png", normalize=True)

